#include <iostream>
#include <vector>
#include <random>
#include <chrono>
#include <immintrin.h>
#include <cmath>

const int NUM_VECTORS = 1000000;  // 向量数量
const int DIMENSIONS = 512;      // 向量维度（确保是16的倍数）

// 生成随机数据
void generateData(std::vector<float>& rowData, std::vector<float>& colData) {
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_real_distribution<float> dis(0.0f, 1.0f);
    
    // 生成行存数据
    for (int i = 0; i < NUM_VECTORS * DIMENSIONS; ++i) {
        rowData[i] = dis(gen);
    }
    
    // 从行存数据转换为列存数据
    for (int i = 0; i < NUM_VECTORS; ++i) {
        for (int j = 0; j < DIMENSIONS; ++j) {
            colData[j * NUM_VECTORS + i] = rowData[i * DIMENSIONS + j];
        }
    }
}

// 行存方式：使用AVX-512优化，计算平方和（不计算平方根）
void rowBasedComputationAVX(const std::vector<float>& rowData, 
                           const std::vector<float>& reference,
                           std::vector<float>& results) {
    for (int i = 0; i < NUM_VECTORS; ++i) {
        __m512 sum_vec = _mm512_setzero_ps();
        int baseIndex = i * DIMENSIONS;
        
        // 每次处理16个维度
        for (int j = 0; j < DIMENSIONS; j += 16) {
            // 加载当前向量的16个维度
            __m512 vec_vals = _mm512_loadu_ps(&rowData[baseIndex + j]);
            // 加载参考向量的16个维度
            __m512 ref_vals = _mm512_loadu_ps(&reference[j]);
            // 计算差值
            __m512 diff = _mm512_sub_ps(vec_vals, ref_vals);
            // 计算平方差并累加
            sum_vec = _mm512_fmadd_ps(diff, diff, sum_vec);
        }
        
        // 将AVX寄存器中的结果求和
        results[i] = _mm512_reduce_add_ps(sum_vec);
    }
}

// 列存方式：使用AVX-512优化，计算平方和（不计算平方根）
void columnBasedComputationAVX(const std::vector<float>& colData,
                              const std::vector<float>& reference,
                              std::vector<float>& results) {
    // 初始化结果数组
    for (int i = 0; i < NUM_VECTORS; ++i) {
        results[i] = 0.0f;
    }
    
    // 对每个维度进行处理
    for (int dim = 0; dim < DIMENSIONS; ++dim) {
        float ref_val = reference[dim];
        __m512 ref_vec = _mm512_set1_ps(ref_val);
        
        // 每次处理16个向量
        for (int i = 0; i < NUM_VECTORS; i += 16) {
            // 加载当前维度的16个值
            __m512 vec_vals = _mm512_loadu_ps(&colData[dim * NUM_VECTORS + i]);
            
            // 计算差值
            __m512 diff = _mm512_sub_ps(vec_vals, ref_vec);
            
            // 计算平方
            __m512 squared = _mm512_mul_ps(diff, diff);
            
            // 加载当前的结果
            __m512 result_vec = _mm512_loadu_ps(&results[i]);
            
            // 累加平方差
            result_vec = _mm512_add_ps(result_vec, squared);
            
            // 存储结果
            _mm512_storeu_ps(&results[i], result_vec);
        }
    }
}

int main() {
    std::cout << "测试配置: " << NUM_VECTORS << " 个向量, " << DIMENSIONS << " 维度" << std::endl;
    
    // 分配内存
    std::vector<float> rowData(NUM_VECTORS * DIMENSIONS);    // 行存数据
    std::vector<float> colData(NUM_VECTORS * DIMENSIONS);    // 列存数据
    std::vector<float> reference(DIMENSIONS);                // 参考向量
    std::vector<float> resultsRow(NUM_VECTORS);              // 行存计算结果
    std::vector<float> resultsCol(NUM_VECTORS);              // 列存计算结果
    
    // 生成随机数据
    generateData(rowData, colData);
    
    // 生成随机参考向量
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_real_distribution<float> dis(0.0f, 1.0f);
    for (int i = 0; i < DIMENSIONS; ++i) {
        reference[i] = dis(gen);
    }
    
    // 测试行存方式性能
    auto start = std::chrono::high_resolution_clock::now();
    rowBasedComputationAVX(rowData, reference, resultsRow);
    auto end = std::chrono::high_resolution_clock::now();
    auto durationRow = std::chrono::duration_cast<std::chrono::microseconds>(end - start).count();
    
    // 测试列存方式性能
    start = std::chrono::high_resolution_clock::now();
    columnBasedComputationAVX(colData, reference, resultsCol);
    end = std::chrono::high_resolution_clock::now();
    auto durationCol = std::chrono::duration_cast<std::chrono::microseconds>(end - start).count();
    
    // 验证结果一致性
    float maxError = 0.0f;
    for (int i = 0; i < NUM_VECTORS; ++i) {
        float error = fabs(resultsRow[i] - resultsCol[i]);
        if (error > maxError) {
            maxError = error;
        }
    }
    
    // 输出结果
    std::cout << "行存方式计算时间: " << durationRow << " 微秒" << std::endl;
    std::cout << "列存方式计算时间: " << durationCol << " 微秒" << std::endl;
    std::cout << "加速比: " << static_cast<float>(durationRow) / durationCol << " 倍" << std::endl;
    
    return 0;
}